#! /usr/bin/env python
# coding=utf-8
# infile = r'test123.fq'

# outfile = r'result111'




from Bio import Seq
from Bio.SeqRecord import SeqRecord

import os
import shutil
import argparse
import sys
import re

# 使用方法 
# conver04_add_gap2simulation.py  -r rawfile -s simulationfile 

parser = argparse.ArgumentParser(
    description=''' 将minibar.py生成的 分配好sample id 的fastq 文件 按每个 amplicon 一个文件进行提取 并自动翻转
    用法:
    barcode_assign_to_file.py  -i result.fq -o result
    由大天才于2022年2月17日创建于浙江农业大学''')

parser.add_argument('-i',
                help='minibar.py输出的分配好的fastq文件')

parser.add_argument('-o',
                help='输出的目录')

args = parser.parse_args()


if not args.i or not args.o:
    parser.print_help()
    sys.exit()




infile = args.i

outfile = args.o



try:
    shutil.rmtree(outfile)
except:
    pass
try:
    os.mkdir(outfile)
except:
    pass



match_reverse = re.compile(r'H[-].+H[+].+')
match_item = re.compile(r'H[-|+].+H[-|+].+')



k = 0
n = 0
t = 0

def write_an_item(l1,l2,l3,l4):
    d = match_item.findall(l1)

    if d!=[] and len(d[0].split(' '))==2:
        
        group = d[0].split(' ')[1]

        tmp_id= d[0].split(' ')[0]


        if re.findall(match_reverse,str(l1))!=[]:
            
            # 需要反转
            l2 = str(SeqRecord(l2).reverse_complement())
            l4 = l4[::-1]
            


        if os.path.exists(outfile+'/'+group+'.fastq'):
            h = open(outfile+'/'+group+'.fastq','a')
            h.write(l1+'\n')
            h.write(l2+'\n')
            h.write(l3+'\n')
            h.write(l4+'\n')
            h.close()
        else:
            h = open(outfile+'/'+group+'.fastq','w')
            h.write(l1+'\n')
            h.write(l2+'\n')
            h.write(l3+'\n')
            h.write(l4+'\n')
            h.close()
  

with open(infile) as fila:
    for i in fila:
        
        i = i.strip()

        if n ==0:
            if t ==1:
                #print(l1,l2,l3,l4)
                write_an_item(l1,l2,l3,l4)
            l1=None
            l2=None
            l3=None
            l4=None

            n = 1
            t = 1    
            k += 1

        if k%10000==0:
            print(k)

        # reads the first line of the fastq
        if n == 1:
            #print(i)
            if len(i)<=1 or i[0] != '@':
                raise BaseException('fastq file error at %s' % i)
            l1 = i

            if len(l1)>999:
                t = 0 

            n += 1
            continue

        if n == 2:
            l2 = i
            n += 1 
            continue

        if n == 3:
            l3 = i
            n += 1 
            continue

        if n == 4:

            l4 = i
            if len(l4) != len(l2):
                t = 0
            else:
                t = 1

            n = 0 
            continue
                

if t ==1:
    write_an_item(l1,l2,l3,l4)



        

        
